import streamlit as st
import akshare as ak
import pandas as pd
import numpy as np
import plotly.graph_objects as go
from plotly.subplots import make_subplots
import random
import time
from stock import stock_hist_em

st.set_page_config(layout="wide")


def fetch_stock_data(stock_code, start_date, end_date):
    try:
        time.sleep(random.uniform(1,2))
        stock_data=stock_hist_em.get_kline_data_sina(code=stock_code,scale=240,limit=240)
        stock_data.set_index('date', inplace=True)
        return stock_data
    except Exception as e:
        st.error(f"获取股票数据时出错: {str(e)}")
        return None


def identify_extreme_points(data, window=5):
    data['high_point'] = data['high'].rolling(window=window, center=True).apply(lambda x: x.argmax() == window // 2)
    data['low_point'] = data['low'].rolling(window=window, center=True).apply(lambda x: x.argmin() == window // 2)
    # Ensure boolean type
    data['high_point'] = data['high_point'].astype(bool)
    data['low_point'] = data['low_point'].astype(bool)
    return data


def connect_extreme_points(data):
    extreme_points = data[data['high_point'] | data['low_point']].copy()
    extreme_points['point_type'] = np.where(extreme_points['high_point'], 'high', 'low')
    extreme_points['value'] = np.where(extreme_points['high_point'], extreme_points['high'], extreme_points['low'])
    return extreme_points


def app():
    st.title("股票K线图与缠论分析")

    col1, col2 = st.columns(2)
    with col1:
        stock_code = st.text_input("输入股票代码（如 sh000001）:", "510300")
    with col2:
        time_period = st.selectbox("选择时间周期", [ "6个月", "1年", "自定义"])

    if time_period == "自定义":
        start_date = st.date_input("选择开始时间:", pd.to_datetime("2024-10-08"))
        end_date = st.date_input("选择结束时间:", pd.to_datetime("2025-01-20"))
    else:
        end_date = pd.Timestamp.now()
        if time_period == "1个月":
            start_date = end_date - pd.DateOffset(months=1)
        elif time_period == "3个月":
            start_date = end_date - pd.DateOffset(months=3)
        elif time_period == "6个月":
            start_date = end_date - pd.DateOffset(months=6)
        else:  # 1年
            start_date = end_date - pd.DateOffset(years=1)

    if st.button("获取数据并分析"):
        with st.spinner("正在获取数据并进行分析..."):
            stock_data = fetch_stock_data(stock_code, start_date.strftime("%Y%m%d"), end_date.strftime("%Y%m%d"))

            # 确保数据按日期升序排列
            stock_data.sort_values('date', inplace=True)
            # 计算真实波幅（True Range）
            stock_data['prev_close'] = stock_data['close'].shift(1)  # 获取前一日收盘价
            stock_data['tr1'] = stock_data['high'] - stock_data['low']  # 当日波幅
            stock_data['tr2'] = abs(stock_data['high'] - stock_data['prev_close'])  # 与前日收盘价的波幅（向上）
            stock_data['tr3'] = abs(stock_data['low'] - stock_data['prev_close'])   # 与前日收盘价的波幅（向下）
            #stock_data['tr'] = stock_data[['tr1', 'tr2', 'tr3']].max(axis=1)  # 三者中的最大值即为真实波幅
            stock_data['tr'] = abs(stock_data['close'] - stock_data['prev_close'])
			# 计算20日ATR
            stock_data['atr20'] = stock_data['tr'].rolling(window=60).mean()
            if stock_data is not None and not stock_data.empty:
                # 缠论分析
                stock_data = identify_extreme_points(stock_data,5)
                extreme_points = connect_extreme_points(stock_data)

                fig = make_subplots(rows=2, cols=1, shared_xaxes=True, vertical_spacing=0.03, row_heights=[0.7, 0.3])

                # K线图
                fig.add_trace(go.Candlestick(
                    x=stock_data.index,
                    open=stock_data['open'],
                    high=stock_data['high'],
                    low=stock_data['low'],
                    close=stock_data['close'],
                    increasing_line_color='red',
                    decreasing_line_color='green',
                    name='K线图'
                ), row=1, col=1)
                fig.add_trace(go.Scatter(x=stock_data.index, y=stock_data['close'].rolling(window=14).mean()-stock_data['atr20'], mode='lines', name='min线', line=dict(color='orange')))
                cutPrice=stock_data['close'].rolling(window=14).mean()[-1]
                cutPrice=cutPrice-stock_data['atr20'][-1]
                fig.add_hline(y=cutPrice, line_dash="dash", line_color="green", annotation_text=f"止损({cutPrice:.2f})")
                fig.add_hline(y=stock_data['close'][-1], line_dash="dash", line_color="red", annotation_text=f"当前({stock_data['close'][-1]:.2f})")

                # 为止损线添加醒目标注
                fig.add_annotation(
                    x=stock_data.index[-1],  # 标注放在最新日期位置
                    y=cutPrice,
                    text=f"止损({cutPrice:.2f})",
                    showarrow=True,
                    arrowhead=2,
                    arrowcolor='green',
                    arrowsize=1,
                    arrowwidth=2,
                    ax=20,  # 箭头水平偏移（正值表示向右）
                    ay=-60,   # 箭头垂直偏移
                    bgcolor="rgba(0,128,0,0.8)",  # 半透明绿色背景
                    bordercolor="green",
                    borderwidth=1,
                    borderpad=4,
                    font=dict(color="white", size=12, family="Arial Black"),  # 白色粗体字体
                    opacity=0.9
                )

                # 计算当前天数-55到-60天内的最高点
                total_days = len(stock_data)
                start_index_1 = max(0, total_days - 60)  # -60天
                end_index_1 = max(0, total_days - 55)    # -55天

                if end_index_1 > start_index_1:
                    period_1 = stock_data.iloc[start_index_1:end_index_1]
                    max_period_1_price = period_1['high'].max()
                    max_period_1_date = period_1['high'].idxmax()
                    
                    # 计算最近-5日到-15天的最高价
                    start_index_2 = max(0, total_days - 15)  # -15天
                    end_index_2 = max(0, total_days - 5)    # -5天
                    
                    if end_index_2 > start_index_2:
                        period_2 = stock_data.iloc[start_index_2:end_index_2]
                        max_period_2_price = period_2['high'].max()
                        max_period_2_date = period_2['high'].idxmax()
                        
                        # 确保起点是较早的日期，终点是较晚的日期
                        if max_period_1_date > max_period_2_date:
                            start_date, start_price = max_period_2_date, max_period_2_price
                            end_date, end_price = max_period_1_date, max_period_1_price
                        else:
                            start_date, start_price = max_period_1_date, max_period_1_price
                            end_date, end_price = max_period_2_date, max_period_2_price
                        
                        # 计算斜率
                        date_diff = (end_date - start_date).days
                        price_diff = end_price - start_price
                        slope = price_diff / date_diff if date_diff != 0 else 0
                        
                        # 创建延伸线的起点和终点（延伸到整个数据范围的两端）
                        first_date = stock_data.index[0]
                        last_date = stock_data.index[-1]
                        
                        # 计算延伸线在两端的价格
                        days_to_first = (first_date - start_date).days
                        price_at_first = start_price + slope * days_to_first
                        
                        days_to_last = (last_date - start_date).days
                        price_at_last = start_price + slope * days_to_last
                        
                        # 画原始线段（实线部分）
                        fig.add_trace(go.Scatter(
                            x=[start_date, end_date],
                            y=[start_price, end_price],
                            mode='lines',
                            line=dict(color='purple', width=3),
                            name='趋势线3'
                        ))
                        
                        # 画向前延伸的虚线（从起点向左延伸）
                        fig.add_trace(go.Scatter(
                            x=[first_date, start_date],
                            y=[price_at_first, start_price],
                            mode='lines',
                            line=dict(color='purple', width=2, dash='dash'),
                            name='趋势线3延伸',
                            showlegend=False
                        ))
                        
                        # 画向后延伸的虚线（从终点向右延伸）
                        fig.add_trace(go.Scatter(
                            x=[end_date, last_date],
                            y=[end_price, price_at_last],
                            mode='lines',
                            line=dict(color='purple', width=2, dash='dash'),
                            name='趋势线3延伸',
                            showlegend=False
                        ))
                        
                        # 添加标记点
                        fig.add_trace(go.Scatter(
                            x=[start_date, end_date],
                            y=[start_price, end_price],
                            mode='markers',
                            marker=dict(size=8, color='purple'),
                            name='趋势关键点3',
                            hovertemplate='日期: %{x}<br>价格: %{y:.2f}<extra></extra>',
                            showlegend=False
                        ))
                        
                        # 添加标注
                        fig.add_annotation(
                            x=start_date, 
                            y=start_price,
                            text=f"-60~-55天高点({start_price:.2f})",
                            showarrow=True,
                            arrowhead=2,
                            arrowcolor='purple',
                            arrowsize=1,
                            arrowwidth=2,
                            ax=0,
                            ay=-60,
                            bgcolor="rgba(128,0,128,0.8)",  # 半透明紫色背景
                            bordercolor="purple",
                            borderwidth=1,
                            borderpad=4,
                            font=dict(color="white", size=12, family="Arial Black"),
                            opacity=0.9
                        )
                        
                        fig.add_annotation(
                            x=end_date, 
                            y=end_price,
                            text=f"-15~-5天高点({end_price:.2f})",
                            showarrow=True,
                            arrowhead=2,
                            arrowcolor='purple',
                            arrowsize=1,
                            arrowwidth=2,
                            ax=0,
                            ay=-60,
                            bgcolor="rgba(128,0,128,0.8)",  # 半透明紫色背景
                            bordercolor="purple",
                            borderwidth=1,
                            borderpad=4,
                            font=dict(color="white", size=12, family="Arial Black"),
                            opacity=0.9
                        )
                
                # 计算60日内最低点
                recent_data = stock_data.tail(60)  # 取最近60个交易日
                min_60_price = recent_data['low'].min()
                min_60_date = recent_data['low'].idxmin()

                # 计算当前天数-5到-26天内的最低点
                total_days = len(stock_data)
                start_index = max(0, total_days - (17+3))  # -26天
                end_index = max(0, total_days - 3)    # -5天

                if end_index > start_index:
                    recent_low_period = stock_data.iloc[start_index:end_index]
                    min_recent_price = recent_low_period['low'].min()
                    min_recent_date = recent_low_period['low'].idxmin()
                    
                    # 确保起点是60日内最低点，终点是-5到-26天内最低点
                    if min_60_date > min_recent_date:
                        start_date, start_price = min_recent_date, min_recent_price
                        end_date, end_price = min_60_date, min_60_price
                    else:
                        start_date, start_price = min_60_date, min_60_price
                        end_date, end_price = min_recent_date, min_recent_price
                    
                    # 计算斜率
                    date_diff = (end_date - start_date).days
                    price_diff = end_price - start_price
                    slope = price_diff / date_diff if date_diff != 0 else 0
                    
                    # 创建延伸线的起点和终点（延伸到整个数据范围的两端）
                    first_date = stock_data.index[0]
                    last_date = stock_data.index[-1]
                    
                    # 计算延伸线在两端的价格
                    days_to_first = (first_date - start_date).days
                    price_at_first = start_price + slope * days_to_first
                    
                    days_to_last = (last_date - start_date).days
                    price_at_last = start_price + slope * days_to_last
                    
                    # 画原始线段（实线部分）
                    fig.add_trace(go.Scatter(
                        x=[start_date, end_date],
                        y=[start_price, end_price],
                        mode='lines',
                        line=dict(color='blue', width=3),
                        name='支撑线'
                    ))
                    
                    
                    
                    # 画向后延伸的虚线（从终点向右延伸）
                    fig.add_trace(go.Scatter(
                        x=[end_date, last_date],
                        y=[end_price, price_at_last],
                        mode='lines',
                        line=dict(color='blue', width=2, dash='dash'),
                        name='支撑线延伸',
                        showlegend=False
                    ))
                    
                    # 添加标记点
                    fig.add_trace(go.Scatter(
                        x=[start_date, end_date],
                        y=[start_price, end_price],
                        mode='markers',
                        marker=dict(size=8, color='blue'),
                        name='关键点位',
                        hovertemplate='日期: %{x}<br>价格: %{y:.2f}<extra></extra>',
                        showlegend=False
                    ))
                    
                    # 添加更清晰的标注
                    fig.add_annotation(
                        x=start_date, 
                        y=start_price,
                        text=f"60日最低({start_price:.2f})",
                        showarrow=True,
                        arrowhead=2,
                        arrowcolor='blue',
                        arrowsize=1,
                        arrowwidth=2,
                        ax=0,  # 箭头水平偏移
                        ay=60,  # 箭头垂直偏移（负值表示向上）
                        bgcolor="rgba(0,0,255,0.8)",  # 半透明蓝色背景
                        bordercolor="blue",
                        borderwidth=1,
                        borderpad=4,
                        font=dict(color="white", size=12, family="Arial Black"),  # 白色粗体字体
                        opacity=0.9
                    )

                    fig.add_annotation(
                        x=end_date, 
                        y=end_price,
                        text=f"近期低点({end_price:.2f})",
                        showarrow=True,
                        arrowhead=2,
                        arrowcolor='blue',
                        arrowsize=1,
                        arrowwidth=2,
                        ax=0,  # 箭头水平偏移
                        ay=60,  # 箭头垂直偏移（正值表示向下）
                        bgcolor="rgba(0,0,255,0.8)",  # 半透明蓝色背景
                        bordercolor="blue",
                        borderwidth=1,
                        borderpad=4,
                        font=dict(color="white", size=12, family="Arial Black"),  # 白色粗体字体
                        opacity=0.9
                    )
                
                # 成交量
                fig.add_trace(go.Bar(
                    x=stock_data.index,
                    y=stock_data['volume'],
                    name='成交量',
                    marker_color='rgba(0, 0, 255, 0.5)'
                ), row=2, col=1)

                # 移动平均线
                #ma_periods = [5, 10, 20]
                #colors = ['blue', 'orange', 'green']

                # for period, color in zip(ma_periods, colors):
                #     ma = stock_data['close'].rolling(window=period).mean()
                #     fig.add_trace(go.Scatter(
                #         x=stock_data.index,
                #         y=ma,
                #         name=f'MA{period}',
                #         line=dict(color=color, width=1)
                #     ), row=1, col=1)

                # 缠论分析：标记最高点和最低点
                fig.add_trace(go.Scatter(
                    x=extreme_points[extreme_points['point_type'] == 'high'].index,
                    y=extreme_points[extreme_points['point_type'] == 'high']['value'],
                    mode='markers',
                    marker=dict(symbol='triangle-up', size=10, color='red'),
                    name='最高点'
                ), row=1, col=1)

                fig.add_trace(go.Scatter(
                    x=extreme_points[extreme_points['point_type'] == 'low'].index,
                    y=extreme_points[extreme_points['point_type'] == 'low']['value'],
                    mode='markers',
                    marker=dict(symbol='triangle-down', size=10, color='green'),
                    name='最低点'
                ), row=1, col=1)

                # 连接最高点和最低点
                for i in range(1, len(extreme_points)):
                    fig.add_trace(go.Scatter(
                        x=[extreme_points.index[i - 1], extreme_points.index[i]],
                        y=[extreme_points['value'].iloc[i - 1], extreme_points['value'].iloc[i]],
                        mode='lines',
                        line=dict(color='purple', width=1),
                        showlegend=False
                    ), row=1, col=1)

                fig.update_layout(
                    title=f"{stock_code} K线图与缠论分析",
                    xaxis_title='日期',
                    yaxis_title='价格',
                    xaxis_rangeslider_visible=False,
                    height=800
                )

                st.plotly_chart(fig, use_container_width=True)

                #st.subheader("缠论分析结果")
                #st.write(f"识别到的极值点数量: {len(extreme_points)}")
                #st.write("极值点详情:")
                #st.dataframe(extreme_points[['point_type', 'value']])

                #st.dataframe(stock_data)
            else:
                st.error("无法获取或处理股票数据，请检查股票代码是否正确，并确保选择了有效的日期范围。")


if __name__ == "__main__":
    app()